#include "ad_construct.h"
#include <cstdio>
#include <cstring>
#include <algorithm>

bool USE_SUFFIX_ARRAY = true;

int getLength[MAX_WORD_LENGTH];
int getY[MAX_WORD_LENGTH];
int getRoot[MAX_WORD_LENGTH];
bool isBadLength[MAX_WORD_LENGTH];
int root_badness_delta[MAX_WORD_LENGTH]; // The change of number of forbidden subwords.

inline int get_length(int plen, const input_data & INPUT){
  return (2 * plen * INPUT.beta.num + INPUT.beta.p + 2 * INPUT.beta.denum - 1) / (2 * INPUT.beta.denum);
}

inline int get_y(int x, const input_data & INPUT){
  return ((2 * INPUT.beta.denum - INPUT.beta.num) * x - INPUT.beta.p) / (INPUT.beta.num - INPUT.beta.denum);
}

inline int get_root(int x, const input_data & INPUT){
  return ((x - 1) * INPUT.beta.denum - INPUT.beta.p) / (INPUT.beta.num - INPUT.beta.denum) + 1;
}

void precalc_arrays(const input_data & INPUT){
  for (int len = 1; len < MAX_WORD_LENGTH; len++){
    getLength[len] = get_length(len, INPUT);
    getY[len] = get_y(len, INPUT);
    getRoot[len] = get_root(len, INPUT);
  }
  int k = INPUT.ALPH_SIZE;
  memset(isBadLength, false, sizeof(isBadLength));
  if (INPUT.beta == ext_rational(k, k-1, 1)){
    isBadLength[3*k] = true;
    isBadLength[4*k] = true;
    int l = k / 2 * k;
    int r = l + k;
    int lq, rq;
    if (k % 2 == 1) lq = rq = (k-1)*(k+1)/2;
    else{
      lq = (k-1)*(k+2)/2;
      rq = lq + 1;
    }
    while (l+1 < lq || rq < r){
      for (int p=l+1; p<lq; p++) isBadLength[p] = true;
      for (int p=rq+1; p<r; p++) isBadLength[p] = true;
      l += k;
      r += k;
      lq += k - 1;
      rq += k + 1;
    }
  }
}

int suffix_array[MAX_DFS_DEPTH][MAX_DFS_DEPTH];
int lcp[MAX_DFS_DEPTH][MAX_DFS_DEPTH];
int leaders_pos[MAX_DFS_DEPTH][MAX_DFS_DEPTH];
int leaders_num[MAX_DFS_DEPTH];

int sa_stack[MAX_DFS_DEPTH];

int qxx1, qxx2;

void rebuild_suffix_array(const char * cword, int cwlen){
  leaders_num[cwlen] = 0;
  int nsz = 0, st_size = 0;
  int r, v, w, i, vlcp;
  suffix_array[cwlen][0] = cwlen;
  lcp[cwlen][0] = 0;
  sa_stack[++st_size] = cwlen;
  char ch = cword[cwlen - 1];
  bool step2 = false;
  if (ch == 'a'){
    suffix_array[cwlen][++nsz] = cwlen - 1;
    lcp[cwlen][nsz] = 1;
    leaders_pos[cwlen][leaders_num[cwlen]++] = nsz;
  }
  else{
    sa_stack[++st_size] = cwlen - 1;
    step2 = true;
  }
  i = 1;
  while (i < cwlen){
    v = suffix_array[cwlen - 1][i];
    vlcp = lcp[cwlen - 1][i];
    r = cwlen - 1 - sa_stack[st_size];
    if (!step2){
      if (lcp[cwlen][nsz] < r){
        lcp[cwlen][nsz + 1] = lcp[cwlen][nsz];
        lcp[cwlen][nsz] = r;
        suffix_array[cwlen][++nsz] = sa_stack[st_size--];
        continue;
      }
      else if (lcp[cwlen][nsz] == r && ch <= cword[v + r]){
        lcp[cwlen][++nsz] = r;
        if (ch == cword[v + r]){
          lcp[cwlen][nsz]++;
          leaders_pos[cwlen][leaders_num[cwlen]++] = nsz;
        }
        suffix_array[cwlen][nsz] = sa_stack[st_size--];
        continue;
      }
    }
    // Step 2
    if (vlcp < cwlen-1-v){
      suffix_array[cwlen][++nsz] = v;
      lcp[cwlen][nsz] = vlcp;
      step2 = false;
    }
    else{
      w = suffix_array[cwlen - 1][i + 1];
      if (ch <= cword[w + vlcp]){
        suffix_array[cwlen][++nsz] = v;
        lcp[cwlen][nsz] = vlcp;
        if (ch == cword[w + vlcp]){
          lcp[cwlen][nsz]++;
          leaders_pos[cwlen][leaders_num[cwlen]++] = nsz;
        }
        step2 = false;
      }
      else{
        sa_stack[++st_size] = v;
        step2 = true;
      }
    }
    i++;
  }

  while (st_size > 1){
    r = cwlen - 1 - sa_stack[st_size];
    suffix_array[cwlen][++nsz] = sa_stack[st_size--];
    lcp[cwlen][nsz - 1] = r;
  }
  lcp[cwlen][cwlen] = -1;

  return;
}

inline int get_lcp_dumb(const char * cword, int clen, int i, int j){
  int res = 0;
  while (i < clen && j < clen && cword[i]==cword[j]){
    i++;
    j++;
    res++;
  }
  return res;
}

int suffs[MAX_ALPH_SIZE][MAX_DFS_DEPTH];
int ssz[MAX_ALPH_SIZE];

int bad_segments_num[MAX_DFS_DEPTH];
int bad_segments_beg[MAX_DFS_DEPTH][MAX_DFS_DEPTH];
int bad_segments_end[MAX_DFS_DEPTH][MAX_DFS_DEPTH];

void handle_root_dumb(char* cword, int clen, char * cpans, int cplen, int badness, dfa * DFA, input_data & INPUT){
  int alph_size = INPUT.ALPH_SIZE;
  bool is_bad = false;
  bad_segments_num[clen] = 0;

  int xlen = getLength[clen] - clen;
  int i, j, s, clcp, ny, cy, rlen;
  char ch;
  for (i = 0; i < alph_size; i++){
    ssz[i] = 0;
  }
  for (i = 1; i <= clen; i++){
    ch = cword[clen-i]-'a';
    suffs[ch][ssz[ch]++] = clen-i;
  }
  for (i = 0; i < xlen; i++){
    ch = cword[i] - 'a';
    for (j = 0; j < ssz[ch]; j++){
      qxx1++;
      s = suffs[ch][j];
      if (s <= xlen) break;
      clcp = 1 + get_lcp_dumb(cword, clen, i + 1, s + 1);
      if (clcp > xlen - i) clcp = xlen - i;
      rlen = s + clcp + getY[clcp] - i;
      if (rlen >= clen){
        bad_segments_beg[clen][bad_segments_num[clen]] = clen + 1;
        bad_segments_end[clen][bad_segments_num[clen]++] = rlen;
        return;
      }
    }
  }

  if (is_bad) return;

  for (i = 0; i < clen && !is_bad; i++){
    ch = cword[i] - 'a';
    for (j = 0; j < ssz[ch]; j++){
      qxx2++;
      s = suffs[ch][j];
      if (i >= s) continue;
      clcp = get_lcp_dumb(cword, clen, i, s);
      if (clcp < clen-s) continue;
      clcp += get_lcp_dumb(cword, clen, 0, i+clcp);
      if (clcp > s - i) clcp = s - i;
      ny = getY[clcp];
      cy = s - i - clcp;
      if (ny >= cy){
        is_bad = true;
        return;
      }
      if (i < xlen){
        if (clcp > xlen - i){
          clcp = xlen - i;
          ny = getY[clcp];
        }
        cy = i + (clen - s) - clcp;
        if (ny >= cy){
          is_bad = true;
          return;
        }
      }
    }
  }

  if (is_bad) return;

  if (!is_bad){
    int blen = clen, cpos = 0;
    int bplen = cplen;
    int nlen = clen + xlen;
    while (blen < nlen){
      char ch = cword[cpos];
      cword[blen] = ch;
      if (ch == cword[blen - alph_size + 1]){
        cpans[bplen] = 'a';
      }
      else{
        cpans[bplen] = 'b';
      }
      blen++;
      bplen++;
      cpos++;
      if (cpos >= clen) cpos = 0;
    }

    cword[blen] = 0;
    /*int bad_pref_num = has_bad_prefix_dumb(cword, blen-1, clen, INPUT);
    if (bad_pref_num > 0){
      return;
    }
    bad_pref_num = has_bad_prefix_dumb(cword+1, blen-1, blen-2, INPUT);
    if (bad_pref_num > 0){
      return;
    }*/

    cpans[bplen] = 0;
    if (blen > MAX_WORD_LENGTH){
      INPUT.dump("Too small maximal length of forbidden word!\n");
      exit(1);
    }
    DFA->add_word(cpans);
    cpans[cplen] = 0;
  }
}

void handle_root(char* cword, int clen, char * cpans, int cplen, int badness, dfa * DFA, input_data & INPUT){
  int alph_size = INPUT.ALPH_SIZE;

  // check that no new forbidden word can appear somewhere
  bool is_bad = (badness > 0);
  int xlen = getLength[clen] - clen;
  int l, i, j;

  bad_segments_num[clen] = 0;

  for (l = 0; l < leaders_num[clen]; l++){
    int pos = leaders_pos[clen][l];
    i = clen - suffix_array[clen][pos];
    int clcp = lcp[clen][pos];
    while (clcp == i){
      j = suffix_array[clen][pos + 1];
      // Some suffix encounters inside XY
      int llen = getRoot[j + i];
      int rlen = clen + getY[i] - j;
      is_bad |= (llen <= clen && clen <= rlen);
      if (llen <= clen) llen = clen + 1;
      if (llen <= rlen){
        bad_segments_beg[clen][bad_segments_num[clen]] = llen;
        bad_segments_end[clen][bad_segments_num[clen]++] = rlen;
      }
      pos++;
      if (clcp > lcp[clen][pos]) clcp = lcp[clen][pos];
    }
  }

  if (is_bad) return;

  // Special cases when Z in a forbidden subword ZTZ of XYX is "cut" into two parts
  // They actually encounter rarely enough
  for (l = 0; l < leaders_num[clen]; l++){
    int pos = leaders_pos[clen][l];
    i = clen - suffix_array[clen][pos];
    int clcp = lcp[clen][pos];
    while (clcp == i){
      j = suffix_array[clen][pos + 1];
      // Y = p + Z1, X = Z2 + T + Z1 + Z2 + q
      if (j + i < xlen && !is_bad && cword[0] == cword[j + i]){
        int add_i = 1 + get_lcp_dumb(cword, clen, 1, j + i + 1);
        if (add_i > j) add_i = j;
        if (add_i > xlen-j-i) add_i = xlen-j-i;
        if (add_i > 0){
          int ny = getY[i + add_i];
          if (ny >= j - add_i){
            is_bad = true;
          }
        }
      }
      // XY = p + Z1 + Z2 + T + Z1, X = Z2 + q
      if (j + i <= clen - i && !is_bad && cword[0] == cword[j + i]){
        int add_i = 1 + get_lcp_dumb(cword, clen, 1, j + i + 1);
        if (add_i > xlen) add_i = xlen;
        if (add_i > (clen-i) - (j+i)) add_i = (clen-i) - (j+i);
        if (add_i > 0){
          int ny = getY[i + add_i];
          if (ny >= (clen-i) - (j+i+add_i)){
            is_bad = true;
          }
        }
      }
      pos++;
      if (clcp > lcp[clen][pos]) clcp = lcp[clen][pos];
    }
  }


  if (!is_bad){
    // construct word here and add it to the automaton
    int blen = clen, cpos = 0;
    int bplen = cplen;
    int nlen = clen + xlen;
    while (blen < nlen){
      char ch = cword[cpos];
      cword[blen] = ch;
      if (ch == cword[blen - alph_size + 1]){
        cpans[bplen] = 'a';
      }
      else{
        cpans[bplen] = 'b';
      }
      blen++;
      bplen++;
      cpos++;
      if (cpos >= clen) cpos = 0;
    }

    cword[blen] = 0;
    /*int bad_pref_num = has_bad_prefix_dumb(cword, blen-1, clen, INPUT);
    if (bad_pref_num > 0){
      return;
    }
    bad_pref_num = has_bad_prefix_dumb(cword+1, blen-1, blen-2, INPUT);
    if (bad_pref_num > 0){
      return;
    }*/

    cpans[bplen] = 0;
    if (blen > MAX_WORD_LENGTH){
      INPUT.dump("Too small maximal length of forbidden word!\n");
      exit(1);
    }
    DFA->add_word(cpans);
    cpans[cplen] = 0;
  }

}

int has_bad_prefix_dumb(const char* cword, int clen, int slen, input_data & INPUT){
  int pos, rlen, i, xlen;
  for (pos=slen; pos<clen; pos++){
    for (rlen = 1; getLength[rlen] <= (pos + 1); rlen++){
      xlen = getLength[rlen];
      for (i=rlen; i<xlen; i++){
        if (cword[pos-i] != cword[pos-i+rlen]) break;
      }
      if (i == xlen) return pos;
    }
  }
  return 0;
}

bool can_append_letter_dumb(const char * cword, int clen, char letter, input_data & INPUT){
  int xlen = getLength[clen] - clen;
  int alph_size = INPUT.ALPH_SIZE;
  for (int i = 0; i < alph_size; i++){
    ssz[i] = 0;
  }
  for (int i = 2; i <= xlen; i++){
    char ch = cword[clen-i]-'a';
    suffs[ch][ssz[ch]++] = clen-i;
  }

  for (int i = clen-3; i >= 0; i--){
    char ch = cword[i]-'a';
    for (int j = 0; j < ssz[ch]; j++){
      int s = suffs[ch][j];
      if (i >= s) continue;
      int clcp = get_lcp_dumb(cword, clen, i, s);
      if (clcp == clen-s && cword[i+clcp] == letter && INPUT.beta <= ext_rational(clen - i + 1, s - i, 0)){
        return false;
      }
    }
  }
  return true;
}


bool can_append_letter(const char * cword, int clen, char letter, input_data & INPUT){
  // We don't check XYX with |X|=1 at the end, because of Pansiot encoding
  int l, i, j;
  for (l = 0; l < leaders_num[clen]; l++){
    int pos = leaders_pos[clen][l];
    i = clen - suffix_array[clen][pos];
    if (i == 1) continue; // Because 00 and 111 are forbidden, |X|=2 sucks
    int clcp = lcp[clen][pos];
    while (clcp == i){
      j = suffix_array[clen][pos + 1];
      if (j + i < clen - i && cword[j + i] == letter && INPUT.beta <= ext_rational(clen - j + 1, clen - i - j, 0)){
        return false;
      }
      pos++;
      if (clcp > lcp[clen][pos]) clcp = lcp[clen][pos];
    }
  }
  return true;
}

inline void do_badness_change(int clen, int sgn){
  int i;
  for (i = 0; i < bad_segments_num[clen]; i++){
    root_badness_delta[ bad_segments_beg[clen][i] ] += sgn;
    root_badness_delta[ bad_segments_end[clen][i] + 1 ] -= sgn;
  }
}

void run_dfs(char * cword, int cwlen, char * cpans, int cplen, int badness, dfa * DFA, input_data & INPUT){
  int alph_size = INPUT.ALPH_SIZE;
  
  cword[cwlen] = 0;
  cpans[cplen] = 0;
  bad_segments_num[cwlen] = 0;
  if (cwlen > alph_size - 1 && USE_SUFFIX_ARRAY){
    rebuild_suffix_array(cword, cwlen);
  }
  if (USE_SUFFIX_ARRAY){
    handle_root(cword, cwlen, cpans, cplen, badness, DFA, INPUT);
  }
  else if (badness == 0 && !isBadLength[cwlen]){
    handle_root_dumb(cword, cwlen, cpans, cplen, badness, DFA, INPUT);
  }
  cword[cwlen] = 0;
  cpans[cplen] = 0;

  if (cwlen == INPUT.ROOT_LENGTH){
    return;
  }

  // Forbid some of the root lengths
  do_badness_change(cwlen, 1);
  
  // Pansiot - choose letter 0 or letter 1
  char ch;
  // Try Pansiot encoding 0
  if (cplen == 0 || cpans[cplen - 1] != 'a'){
    ch = cword[cwlen - alph_size + 1];
    bool can = USE_SUFFIX_ARRAY ? can_append_letter(cword, cwlen, ch, INPUT) : can_append_letter_dumb(cword, cwlen, ch, INPUT);
    if (can){
      cword[cwlen] = ch;
      cpans[cplen] = 'a';
      run_dfs(cword, cwlen + 1, cpans, cplen + 1, badness + root_badness_delta[cwlen + 1], DFA, INPUT);
    }
  }

  // Try Pansiot encoding 1
  if (cplen <= 1 || cpans[cplen-1] != 'b' ||cpans[cplen-2] != 'b'){
    int zback = (cplen == 0) ? -1 : (cpans[cplen-1] == 'a') ? alph_size + 1 : alph_size;
    ch = (cwlen >= zback) ? cword[cwlen - zback] : ('a' + alph_size - 1);
    bool can = USE_SUFFIX_ARRAY ? can_append_letter(cword, cwlen, ch, INPUT) : can_append_letter_dumb(cword, cwlen, ch, INPUT);
    if (can){
      cword[cwlen] = ch;
      cpans[cplen] = 'b';
      run_dfs(cword, cwlen + 1, cpans, cplen + 1, badness + root_badness_delta[cwlen + 1], DFA, INPUT);
    }
  }
  do_badness_change(cwlen, -1);
}

i64 construct_antidictionary(grate_config & CONFIG, input_data & INPUT, words_queue * WORDS_QUEUE, dfa * DFA){

  precalc_arrays(INPUT);

  int alph_size = INPUT.ALPH_SIZE;

  char cword[MAX_WORD_LENGTH];
  int cwlen = 0;
  char cpans[MAX_WORD_LENGTH];
  int cplen = 0;
  for (int i = 0; i < alph_size - 1; i++){
    cword[cwlen++] = 'a' + i;
  }
  cword[cwlen] = 0;

  suffix_array[alph_size - 1][0] = alph_size - 1;
  lcp[alph_size - 1][0] = 0;
  for (int i = 0; i < alph_size - 1; i++){
    suffix_array[alph_size - 1][i + 1] = i;
    lcp[alph_size - 1][i + 1] = 0;
  }
  
  qxx1 = 0;
  qxx2 = 0;
  run_dfs(cword, cwlen, cpans, cplen, 0, DFA, INPUT);
  //fprintf(stderr, "%d %d\n", qxx1, qxx2);

  delete WORDS_QUEUE;
  return 0;
}
